import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.ticker import MultipleLocator


params = {
    "axes.labelsize": 22,
    "font.size": 22,
    "legend.fontsize": 20,
    "xtick.labelsize": 20,
    "ytick.labelsize": 20,
    "grid.color": "k",
    "grid.linestyle": ":",
    "grid.linewidth": 1,
}
mpl.rcParams.update(params)

estimate = pd.read_csv("international_migration_flow.csv", na_filter=False)
df = estimate.groupby("migration_month", as_index=False).agg({"num_migrants": "sum"})
df = df.sort_values(by="migration_month")
moving_avg = []
v = list(df.num_migrants)
for i in range(len(df)):
    if i < 1:
        r = np.mean(v[: (i + 2)])
    else:
        r = np.mean(v[(i - 1) : (i + 2)])
    moving_avg.append(r)

fig, ax = plt.subplots(figsize=(22, 4))
plt.plot(
    list(range(len(df))),
    list(df["num_migrants"] / 1e6),
    "--",
    linewidth=1,
    color="gray",
    label="Monthly",
)
plt.plot(
    list(range(len(df))),
    list(np.array(moving_avg) / 1e6),
    "-",
    color="k",
    label="Quarter moving average",
)

plt.ylabel("# Migrants (Millions)")
plt.ticklabel_format(style="sci", axis="y", scilimits=(2, 0))
plt.xlim(0, len(df) - 1)
ax.xaxis.set_major_locator(MultipleLocator(3))
ax.xaxis.set_minor_locator(MultipleLocator(1))
ax.tick_params(which="major", length=6)
ax.tick_params(which="minor", length=3)
ax.set_xticklabels(
    [
        "",
        "Jan 2019",
        "Apr 2019",
        "Jul 2019",
        "Oct 2019",
        "Jan 2020",
        "Apr 2020",
        "Jul 2020",
        "Oct 2020",
        "Jan 2021",
        "Apr 2021",
        "Jul 2021",
        "Oct 2021",
        "Jan 2022",
        "Apr 2022",
        "Jul 2022",
        "Oct 2022",
    ],
    minor=False,
    fontsize=13,
)
plt.gca().yaxis.grid(True, linestyle="-", c="lightgray")
plt.legend(loc=4, fontsize=13)
plt.savefig(
    "figures/fig1.pdf",
    facecolor="white",
    transparent=False,
    bbox_inches="tight",
    dpi=300,
)
